package main

import (
    "sync/atomic"
    "time"
    "math/rand"
    "sync"
    "io"
    "net"
)

type TcpProxy struct {
    clientConn  *net.TCPConn
    serverConn  *net.TCPConn
    clientAddr  string
    serverAddr  string
    closed      int32
    err         error
}

func newTcpProxy(client *net.TCPConn) *TcpProxy {
    p := &TcpProxy{}
    p.clientConn = client
    p.clientAddr = client.RemoteAddr().String()
    p.closed = 0
    p.err = nil
    return p
}

func (p *TcpProxy) connectToServer(addr string) error {
    tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
    if err != nil {
        return err
    }
    p.serverConn, err = net.DialTCP("tcp", nil, tcpAddr)
    if err != nil {
        return err
    }
    
    p.serverAddr = addr
    return nil
}

func (p *TcpProxy) runProxyAndWait() {
    ioCopy := func(dst *net.TCPConn, src *net.TCPConn, wg *sync.WaitGroup) {
        defer wg.Done()
        _, err := io.Copy(dst, src)
        if p.closed == 1 {
            return
        }
        p.close(err)
    }
    
    wg := sync.WaitGroup{}
    wg.Add(2)
    go ioCopy(p.clientConn, p.serverConn, &wg)
    go ioCopy(p.serverConn, p.clientConn, &wg)
    
    wg.Wait()
}

func (p *TcpProxy) close(err error) {
    if !atomic.CompareAndSwapInt32(&p.closed, 0, 1) {
        return
    }
    p.err = err
    if p.serverConn != nil {
        p.serverConn.Close();
    }
    if (p.clientConn != nil) {
        p.clientConn.Close();
    }
}


type TcpProxyService struct {
    config  *TcpProxyConfig
    server  *TcpServer
    proxy   map[string]*TcpProxy
    lock    sync.Mutex
}

func newTcpProxyService(config *TcpProxyConfig) *TcpProxyService {
    s := &TcpProxyService{}
    s.config = config
    s.server = newTcpServer(config.ListenAddr)
    s.server.handler = s
    s.proxy = make(map[string]*TcpProxy)
    return s
}

func (s *TcpProxyService) setConfig(config *TcpProxyConfig) {
    s.config = config
}

func (s *TcpProxyService) service(conn *net.TCPConn) {
    config := s.config
    serverNum := len(config.Server)
    if serverNum == 0 {
        conn.Close()
        return
    }
    
    proxy := newTcpProxy(conn)
    var err error = nil
    
    if config.hashTable == nil {
        r := rand.New(rand.NewSource(time.Now().UnixNano()))
        index := r.Intn(serverNum)
        
        for i := 0; i < serverNum; i++ {
            addr := config.Server[index]
            err = proxy.connectToServer(addr)
            if err == nil {
                break
            }
            index = (index+1)%serverNum
        }
    } else {
        ip, _, _ := net.SplitHostPort(proxy.clientAddr)
        server := config.hashTable.get(ip)
        err = proxy.connectToServer(server)
    }
    
    if err != nil {
        proxy.close(nil)
        debugLog.Printf("tcp proxy connect to server failed. server down. local(%s)\n",
                        s.config.ListenAddr)
        return
    }
    
    s.addProxy(proxy)
    proxy.runProxyAndWait()
    s.delProxy(proxy)
}

func (s *TcpProxyService) addProxy(proxy *TcpProxy) {
    s.lock.Lock()
    defer s.lock.Unlock()
    
    s.proxy[proxy.clientAddr] = proxy
    
    debugLog.Printf("new tcp proxy. client(%s) local(%s) num(%d)\n", 
                    proxy.clientAddr, s.config.ListenAddr, len(s.proxy))
}

func (s *TcpProxyService) delProxy(proxy *TcpProxy) {
    s.lock.Lock()
    defer s.lock.Unlock()
    
    s.proxy[proxy.clientAddr] = nil
    delete(s.proxy, proxy.clientAddr)
    
    debugLog.Printf("del tcp proxy. client(%s) local(%s) num(%d) err(%v)\n", 
                    proxy.clientAddr, s.config.ListenAddr, len(s.proxy), proxy.err)
}

func (s *TcpProxyService) start() error {
    return s.server.start()
}

func (s *TcpProxyService) close() {
    s.lock.Lock()
    for _, p := range s.proxy {
        p.close(nil)
    }
    s.lock.Unlock()
    
    s.server.close()
}

type TcpProxyServiceMgr struct {
    list map[string]*TcpProxyService
    lock sync.Mutex
}

var gTcpProxyServiceMgr *TcpProxyServiceMgr = nil
func TcpProxyServiceMgrInit() {
    gTcpProxyServiceMgr = &TcpProxyServiceMgr{}
    gTcpProxyServiceMgr.list = make(map[string]*TcpProxyService)
}

func (mgr *TcpProxyServiceMgr) startService(config []*TcpProxyConfig) {
    mgr.reloadService(config)
}

func (mgr *TcpProxyServiceMgr) reloadService(config []*TcpProxyConfig) {
    mgr.lock.Lock()
    defer mgr.lock.Unlock()
    
    newProxy := make(map[string]*TcpProxyConfig)
    for _, c := range config {
        newProxy[c.ListenAddr] = c
    }
    
    // add new
    for k, c := range newProxy {
        p, ok := mgr.list[k]
        if !ok {
            p = newTcpProxyService(c)
            err := p.start()
            if err != nil {
                mainLog.Printf("start tcp proxy service failed. listen(%s) err(%v)\n", c.ListenAddr, err)
                continue
            }
            mgr.list[k] = p
        } else {
            p.setConfig(c)
        }
    }
    
    // del old
    for k, p := range mgr.list {
        _, ok := newProxy[k]
        if !ok {
            p.close()
            delete(mgr.list, k)
        }
    }
}